﻿# 赵世钰 迷宫游戏 类似 代替
import numpy as np

from Game.G2FrozenLake.agent import Agent
from Game.G2FrozenLake.env import FrozenLakeEnv


def train_agent(env, agent, steps=1000):
    total_reward = 0
    episode = 0

    while steps > 0:
        # 初始化环境
        current_state = env.get_current_position()["state_index"]
        episode_steps = 0

        while True:
            # 选择动作
            action = agent.choose_action()  # 0,1,2,3 # 行 列 左下右上

            # 执行动作
            _, reward, done, truncated, info = env.step(action)
            next_state = env.get_current_position()["state_index"]

            # 更新 Q 表
            agent.update(current_state, action, reward, next_state)

            # 统计累计奖励
            total_reward += reward
            current_state = next_state
            steps -= 1
            episode_steps += 1

            # 终止条件
            if done or truncated or steps <= 0:
                break

        print(
            f"Episode {episode} | Steps: {episode_steps} | Total Reward: {total_reward:.2f}"
        )
        episode += 1
        env.env.reset()


import matplotlib.pyplot as plt


def plot_results(algorithm_name, rewards, best_arm_rates):
    plt.figure(figsize=(12, 5))

    # 平均奖励曲线
    plt.subplot(1, 2, 1)
    plt.plot(np.cumsum(rewards) / np.arange(1, len(rewards) + 1), label=algorithm_name)
    plt.xlabel("Steps")
    plt.ylabel("Average Reward")
    plt.legend()

    # 最佳臂选择比例
    plt.subplot(1, 2, 2)
    plt.plot(
        np.cumsum(best_arm_rates) / np.arange(1, len(best_arm_rates) + 1),
        label=algorithm_name,
    )
    plt.xlabel("Steps")
    plt.ylabel("Best Arm Rate")
    plt.legend()

    plt.show()


# 创建环境和智能体
env = FrozenLakeEnv()
agent = Agent(env)
train_agent(env, agent, steps=2000)
